Skip to content

[None][fix] AutoDeploy: set enable_spec_decode on ADEngine for disagg#15260

Merged
govind-ramnarayan merged 3 commits into
NVIDIA:mainfrom
Shixiaowei02:fix/ad-disagg-enable-spec-decode
Jun 12, 2026
Merged

[None][fix] AutoDeploy: set enable_spec_decode on ADEngine for disagg#15260
govind-ramnarayan merged 3 commits into
NVIDIA:mainfrom
Shixiaowei02:fix/ad-disagg-enable-spec-decode

Conversation

@Shixiaowei02

@Shixiaowei02 Shixiaowei02 commented Jun 11, 2026

Copy link
Copy Markdown
Collaborator

This pull request makes a small but important change to the ad_executor.py file to improve compatibility between ADEngine and shared PyExecutor code. Specifically, it ensures that the spec-decode flags expected by the shared code are set correctly when ADEngine is used.

  • Compatibility improvements:
    • Set is_spec_decode and enable_spec_decode flags in the ADEngine initializer to match what shared PyExecutor code expects, ensuring proper handling of spec-decode features.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Summary by CodeRabbit

  • Bug Fixes
    • Resolved speculative decoding compatibility issue in AutoDeploy engine by properly initializing required configuration flags during startup.

@Shixiaowei02 Shixiaowei02 requested a review from a team as a code owner June 11, 2026 11:18
@Shixiaowei02 Shixiaowei02 requested a review from galagam June 11, 2026 11:18
ADEngine subclasses the abstract ModelEngine and does not run
PyTorchModelEngine.__init__, so it never set `enable_spec_decode`. After
NVIDIA#14546 added an unguarded `self.model_engine.enable_spec_decode` read in
`_prepare_disagg_gen_transmission_complete` (the disagg generation handoff
path that ADEngine traverses via NVIDIA#14057 AutoDeploy Basic Disagg Support),
AutoDeploy disaggregated runs crash with:
  AttributeError: 'ADEngine' object has no attribute 'enable_spec_decode'

NVIDIA#14546 and NVIDIA#14057 each passed CI independently but conflict semantically
once both are on main. Set `is_spec_decode`/`enable_spec_decode` in
ADEngine.__init__, mirroring PyTorchModelEngine
(enable_spec_decode == spec_config is not None), so ADEngine satisfies the
ModelEngine attribute contract that shared PyExecutor code relies on.

Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
@Shixiaowei02 Shixiaowei02 force-pushed the fix/ad-disagg-enable-spec-decode branch from 835447b to 73ed620 Compare June 11, 2026 11:18
@coderabbitai

coderabbitai Bot commented Jun 11, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: b85bb9e2-0662-4ef0-909a-fddc8e1b89ba

📥 Commits

Reviewing files that changed from the base of the PR and between 835fd61 and 73ed620.

📒 Files selected for processing (1)
  • tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

📝 Walkthrough

Walkthrough

ADEngine.init now initializes spec-decode compatibility flags (is_spec_decode and enable_spec_decode) based on whether spec_config is present, filling fields previously left unset when ADEngine bypasses PyTorchModelEngine.init.

Changes

Spec-decode state initialization

Layer / File(s) Summary
Initialize spec-decode flags
tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
ADEngine.__init__ sets is_spec_decode and enable_spec_decode from spec_config presence to satisfy shared PyExecutor expectations for spec-decode behavior.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly describes the specific change (setting enable_spec_decode on ADEngine) and the context (AutoDeploy disagg), making it clear and specific enough for history scanning.
Description check ✅ Passed The description explains the issue, solution, and includes a checked PR checklist, though Test Coverage section is empty and could be more explicit about test validation.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands and usage tips.

@Shixiaowei02

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53561 [ run ] triggered by Bot. Commit: 73ed620 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53561 [ run ] completed with state FAILURE. Commit: 73ed620
/LLM/main/L0_MergeRequest_PR pipeline #42710 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@Shixiaowei02

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53583 [ run ] triggered by Bot. Commit: 73ed620 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53583 [ run ] completed with state FAILURE. Commit: 73ed620
/LLM/main/L0_MergeRequest_PR pipeline #42730 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@govind-ramnarayan

Copy link
Copy Markdown
Collaborator

Thank you for flagging this and submitting the PR!

@govind-ramnarayan

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53706 [ run ] triggered by Bot. Commit: 73ed620 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53706 [ run ] completed with state FAILURE. Commit: 73ed620
/LLM/main/L0_MergeRequest_PR pipeline #42837 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@govind-ramnarayan

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53719 [ run ] triggered by Bot. Commit: 95b10a6 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53719 [ run ] completed with state FAILURE. Commit: 95b10a6
/LLM/main/L0_MergeRequest_PR pipeline #42848 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@govind-ramnarayan

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53918 [ run ] triggered by Bot. Commit: 68d8f45 Link to invocation

@govind-ramnarayan govind-ramnarayan enabled auto-merge (squash) June 12, 2026 18:32
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #53918 [ run ] completed with state SUCCESS. Commit: 68d8f45
/LLM/main/L0_MergeRequest_PR pipeline #43014 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

@govind-ramnarayan govind-ramnarayan merged commit 19ae053 into NVIDIA:main Jun 12, 2026
8 checks passed
@Shixiaowei02 Shixiaowei02 deleted the fix/ad-disagg-enable-spec-decode branch June 29, 2026 05:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants